"""
从yolov5的检测结果中提取与人有关的边界框信息，包括坐标和检测概率，保存为pkl文件
"""
import os
import pickle

# 直接指定路径
labelPath = r'D:\file\postgrad\experiment\bird_ava_dataset\merge_label_txt'  # YOLOv5检测结果的位置
avaMin_dense_proposals_train_path = r'D:\file\postgrad\experiment\bird_ava_dataset\avaMin_dense_proposals_train.pkl'  # 保存为pkl的路径
avaMin_dense_proposals_val_path = r'D:\file\postgrad\experiment\bird_ava_dataset\avaMin_dense_proposals_val.pkl'  # 保存为pkl的路径
showPkl = True  # 是否可视化

train_dict = {}  # 存储训练集的边界框信息
val_dict = {}  # 存储测试集的边界框信息

files = os.listdir(labelPath)
files.sort(key=lambda arr: (int(arr.split('_')[0]), int(arr.split('_')[1][:-4])))  # 对文件名进行排序


# 遍历文件并按视频号进行分类
for idx, name in enumerate(files):
    temp_file_name = name.split("_")[0]  # 从文件名中提取视频名
    temp_video_ID = name.split("_")[1].split('.')[0]
    temp_video_ID = int(temp_video_ID)
    temp_video_ID = str(temp_video_ID).zfill(4)

    key = temp_file_name + ',' + temp_video_ID  # 视频名，视频ID

    with open(os.path.join(labelPath, name)) as temp_txt:
        temp_data_txt = temp_txt.readlines()
        results = []  # 初始一个空列表，存储当前文件中提取的边界框信息
        for i in temp_data_txt:
            j = i.split(' ')  # 将当前行按空格分割，得到列表，包括检测的类别、坐标和概率信息
            y = j
            y[1] = float(j[1]) - float(j[3]) / 2  # top left x
            y[2] = float(j[2]) - float(j[4]) / 2  # top left y
            y[3] = float(j[1]) + float(j[3])  # bottom right x
            y[4] = float(j[2]) + float(j[4])  # bottom right y

            # 将提取的边界框信息存储到结果列表中
            results.append([y[1], y[2], y[3], y[4]])

        # 每个视频按顺序分配到训练集或测试集
        if int(temp_file_name) % 5 == 0:
            val_dict[key] = results
        else:
            train_dict[key] = results


# 保存 train.pkl
with open(avaMin_dense_proposals_train_path, "wb") as train_pkl_file:
    pickle.dump(train_dict, train_pkl_file)

# 保存 val.pkl
with open(avaMin_dense_proposals_val_path, "wb") as val_pkl_file:
    pickle.dump(val_dict, val_pkl_file)

# 显示 pkl 文件中的内容
if showPkl:
    print("Train data:")
    for i in train_dict:
        print(i, train_dict[i])
    print("Val data:")
    for i in val_dict:
        print(i, val_dict[i])


